package org.gwh.airagknowledge.core.rag;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.gwh.airagknowledge.core.embedding.EmbeddingGenerator;
import org.gwh.airagknowledge.core.llm.impl.LlmFactory;
import org.gwh.airagknowledge.entity.DocumentChunk;
import org.gwh.airagknowledge.entity.QAHistory;
import org.gwh.airagknowledge.entity.VectorEmbedding;
import org.gwh.airagknowledge.repository.QAHistoryRepository;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
@Component
@RequiredArgsConstructor
public class RagQueryEngine {

    private final EmbeddingGenerator embeddingGenerator;
    private final LlmFactory llmFactory;
    private final QAHistoryRepository qaHistoryRepository;

    private static final int DEFAULT_TOP_K = 5;

    /**
     * 统一查询接口 - 根据是否提供知识库ID决定查询方式
     */
    public QAHistory query(Long knowledgeBaseId, String question, double temperature) {
        if (knowledgeBaseId != null && knowledgeBaseId > 0) {
            // 使用知识库进行问答
            return queryWithKnowledgeBase(knowledgeBaseId, question, DEFAULT_TOP_K, temperature);
        } else {
            // 直接问答，不使用知识库
            return queryDirect(question, temperature);
        }
    }

    /**
     * 使用知识库的RAG查询流程
     */
    private QAHistory queryWithKnowledgeBase(Long knowledgeBaseId, String question, int topK, double temperature) {
        try {
            log.info("Processing knowledge base query: {} with temperature: {}", question, temperature);
            
            // 设置LLM温度参数
            llmFactory.setTemperature(temperature);
            
            // 1. 检索相关文档
            List<VectorEmbedding> similarEmbeddings = embeddingGenerator.findSimilarEmbeddings(knowledgeBaseId, question, topK);
            
            // 2. 提取文档内容
            List<DocumentChunk> relevantChunks = similarEmbeddings.stream()
                    .map(VectorEmbedding::getChunk)
                    .collect(Collectors.toList());
            
            // 3. 构建上下文
            String context = buildContext(relevantChunks);
            
            // 4. 生成答案
            String answer = llmFactory.getCurrentLlmService().generateAnswer(question, context);
            
            // 5. 构建源引用
            Map<String, Object> sources = buildSources(relevantChunks);
            
            // 6. 构建元数据
            Map<String, Object> metadata = buildMetadata(temperature, topK, true);
            
            // 7. 保存问答历史
            return saveQAHistory(question, answer, sources, metadata);
        } catch (Exception e) {
            log.error("Error processing knowledge base query: {}", question, e);
            throw new RuntimeException("Failed to process query", e);
        }
    }

    /**
     * 直接问答流程，不使用知识库
     */
    private QAHistory queryDirect(String question, double temperature) {
        try {
            log.info("Processing direct query: {} with temperature: {}", question, temperature);
            
            // 设置LLM温度参数
            llmFactory.setTemperature(temperature);
            
            // 直接使用LLM生成答案，不检索文档
            String answer = llmFactory.getCurrentLlmService().generateAnswer(question, "");
            
            // 构建元数据
            Map<String, Object> metadata = buildMetadata(temperature, 0, false);
            
            // 保存问答历史
            return saveQAHistory(question, answer, Collections.emptyMap(), metadata);
        } catch (Exception e) {
            log.error("Error processing direct query: {}", question, e);
            throw new RuntimeException("Failed to process direct query", e);
        }
    }
    
    /**
     * 构建元数据信息
     */
    private Map<String, Object> buildMetadata(double temperature, int topK, boolean usedKnowledgeBase) {
        Map<String, Object> metadata = new HashMap<>();
        metadata.put("model", llmFactory.getCurrentLlmService().getModelName());
        metadata.put("temperature", temperature);
        metadata.put("used_knowledge_base", usedKnowledgeBase);
        if (usedKnowledgeBase && topK > 0) {
            metadata.put("topK", topK);
        }
        return metadata;
    }
    
    /**
     * 保存问答历史记录
     */
    private QAHistory saveQAHistory(String question, String answer, Map<String, Object> sources, Map<String, Object> metadata) {
        QAHistory qaHistory = QAHistory.builder()
                .question(question)
                .answer(answer)
                .sources(sources)
                .metadata(metadata)
                .build();
        
        return qaHistoryRepository.save(qaHistory);
    }
    
    /**
     * 构建上下文内容
     */
    private String buildContext(List<DocumentChunk> chunks) {
        StringBuilder context = new StringBuilder();
        for (int i = 0; i < chunks.size(); i++) {
            context.append("--- Document ").append(i + 1).append(" ---\n");
            context.append(chunks.get(i).getContent()).append("\n\n");
        }
        return context.toString();
    }
    
    /**
     * 构建源引用信息
     */
    private Map<String, Object> buildSources(List<DocumentChunk> chunks) {
        Map<String, Object> sources = new HashMap<>();
        for (int i = 0; i < chunks.size(); i++) {
            DocumentChunk chunk = chunks.get(i);
            Map<String, Object> sourceInfo = new HashMap<>();
            sourceInfo.put("document_id", chunk.getDocument().getId());
            sourceInfo.put("document_name", chunk.getDocument().getFileName());
            sourceInfo.put("chunk_id", chunk.getId());
            sourceInfo.put("chunk_index", chunk.getChunkIndex());
            sources.put("source_" + (i + 1), sourceInfo);
        }
        return sources;
    }
} 